package com.zyk.leetcode;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author zhangsan
 * @date 2021/5/7 10:24
 */
public class C398 {

    // 给你一个数组, 可能很大, 数组可以有重复值
    // pick: 返回给定数的索引, 索引可能有多个, 等概率返回一个即可
    static class Solution {

        Map<Integer, List<Integer>> indexMap = new HashMap<>();

        public Solution(int[] nums) {
            for (int i = 0; i < nums.length; i++) {
                if (!indexMap.containsKey(nums[i])) {
                    indexMap.put(nums[i], new ArrayList<>());
                }
                indexMap.get(nums[i]).add(i);
            }
        }

        public int pick(int target) {
            List<Integer> list = indexMap.get(target);
            int i = (int) (Math.random() * list.size());
            return list.get(i);
        }

    }

    // for test
    public static void main(String[] args) {
        int[] nums = new int[]{1, 2, 3, 3, 3};
        Solution solution = new Solution(nums);

        // pick(3) 应该返回索引 2,3 或者 4。每个索引的返回概率应该相等。
        System.out.println(solution.pick(3));

        // pick(1) 应该返回 0。因为只有nums[0]等于1。
        System.out.println(solution.pick(1));

    }

}
